import math

import torch
import torch.nn as nn
import os
import csv
import smtplib
from email.mime.text import MIMEText
from email.mime.image import MIMEImage
from email.mime.multipart import MIMEMultipart
import platform
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

# 超参数
batch_size = 8 if platform.system() == 'Windows' else 64
lr = 0.1 if platform.system() == 'Windows' else 0.001
epoch_num = 1 if platform.system() == 'Windows' else 100

# 固定参数
cpu_workers = 6
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

def init(func):
    # net
    net = func().to(device)
    # optimizer
    optimizer = torch.optim.Adam(net.fc.parameters(), lr=lr)
    # loss
    loss_func = nn.CrossEntropyLoss()
    # 学习率调整的方法
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5,
                                                gamma=0.9)  # 采用指数衰减，每5次变成上一次学习率的0.9倍
    return net, optimizer, loss_func, scheduler


# 一些与训练本身无关的工作
data_dir = '../../data'

if not os.path.exists(data_dir):
    os.makedirs(data_dir)
if not os.path.exists('../txt'):
    os.makedirs('../txt')
if not os.path.exists('../csv'):
    os.mkdir('../csv')
if not os.path.exists('../model'):
    os.mkdir('../model')
if not os.path.exists('../img'):
    os.mkdir('../img')
if not os.path.exists('../note'):
    os.mkdir('../note')

# cvs
train_csv_path = f'../csv/train.csv'
train_file = open(train_csv_path, 'a+', encoding='utf-8', newline='')
train_csv_writer = csv.writer(train_file)

test_csv_path = f'../csv/test.csv'
test_file = open(test_csv_path, 'a+', encoding='utf-8', newline='')
test_csv_writer = csv.writer(test_file)

train_accuracy_s = []
test_accuracy_s = []


def my_print(verbose, epoch=None, step=None, size=None, loss=None, accuracy=None, _lr=None, ):
    if verbose == 0:
        print(f'train epoch:{epoch},step:{step}/{size},loss:{loss},accuracy:{accuracy},lr:{_lr}')
    if verbose == 1:
        print(f'==========第{epoch}次训练效果==========')
        print(f'train epoch:{epoch},loss:{loss},accuracy:{accuracy},lr:{_lr}')
        train_accuracy_s.append((accuracy, epoch))
        train_csv_writer.writerow([epoch, loss, accuracy, _lr])
        print(f'=======================================')
    if verbose == 2:
        print(f'==========第{epoch}次测试效果==========')
        print(f'test epoch:{epoch},loss:{loss},accuracy:{accuracy}')
        test_accuracy_s.append((accuracy, epoch))
        test_csv_writer.writerow([epoch, loss, accuracy])
        print(f'=======================================')
    if verbose == 3:
        print(f'==========脚本运行结束==========')
        print('训练集准确率最高的模型: ', sorted(train_accuracy_s, reverse=True)[:10])
        print('测试集准确率最高的模型: ', sorted(test_accuracy_s, reverse=True)[:10])
        print(f'=======================================')


class MyEmail():
    def __init__(self, subject):
        self.host = 'smtp.qq.com'
        self.port = 465
        self.user = 'xueguopeng@foxmail.com'
        self.password = 'islclqzazvnidbdc'
        self.subject = subject

    def send_good_email(self):
        if len(train_accuracy_s) != 0 and len(test_accuracy_s) != 0:
            my_print(verbose=3)
            msg = '<p>训练集准确率最高的模型: ' + str(sorted(train_accuracy_s, reverse=True)[:3]) \
                  + '</p><p>测试集准确率最高的模型: ' + str(sorted(test_accuracy_s, reverse=True)[:3]) + '</p>' \
                  + '<img src="cid:loss_img"><img src="cid:acc_img">'
        else:
            msg = 'train_accuracy_s or test_accuracy_s is None'

        subject = self.subject + '-脚本运行结束通知'

        email_client = smtplib.SMTP_SSL(host=self.host)
        # 设置发件人邮箱的域名和端口
        email_client.connect(host=self.host, port=self.port)

        # 登陆邮件，权限验证，password为邮箱密码
        result = email_client.login(user=self.user, password=self.password)
        print("登录结果", result)

        # 发送邮件，from_addr：发送人，to_addrs：收件人 ，msg：发送的文本
        message = MIMEMultipart()
        message['From'] = self.user
        message['To'] = self.user
        message['Subject'] = subject

        # 发送图片，增加图片标签
        if os.path.exists('../img/loss-epoch.png'):
            with open('../img/loss-epoch.png', 'rb') as f:
                image_data = f.read()
            image = MIMEImage(image_data)
            image.add_header('Content-ID', '<loss_img>')
            message.attach(image)
        if os.path.exists('../img/accuracy-epoch.png'):
            with open('../img/accuracy-epoch.png', 'rb') as f:
                image_data = f.read()
            image = MIMEImage(image_data)
            image.add_header('Content-ID', '<acc_img>')
            message.attach(image)

        message.attach(MIMEText(msg, 'html'))
        email_client.sendmail(from_addr=self.user, to_addrs=self.user, msg=message.as_string())
        # 关闭邮件客户端
        email_client.close()

    def send_bad_email(self, msg):

        subject = self.subject + '-脚本运行失败通知'

        email_client = smtplib.SMTP_SSL(host=self.host)
        # 设置发件人邮箱的域名和端口
        email_client.connect(host=self.host, port=self.port)

        # 登陆邮件，权限验证，password为邮箱密码
        result = email_client.login(user=self.user, password=self.password)
        print("登录结果", result)

        # 发送邮件，from_addr：发送人，to_addrs：收件人 ，msg：发送的文本
        message = MIMEMultipart()
        message['From'] = self.user
        message['To'] = self.user
        message['Subject'] = subject
        message.attach(MIMEText(msg))
        email_client.sendmail(from_addr=self.user, to_addrs=self.user, msg=message.as_string())
        # 关闭邮件客户端
        email_client.close()


def analyze():
    if os.path.exists('../csv/test.csv') and os.path.exists('../csv/train.csv'):
        train_res = pd.read_csv('../csv/train.csv')
        train_epoch = train_res['train_epoch']
        train_loss = train_res['train_loss']
        train_acc = train_res['train_accuracy']
        train_lr = train_res['lr']

        test_res = pd.read_csv('../csv/test.csv')
        test_res = test_res.sort_values(by='test_epoch')
        test_epoch = test_res['test_epoch']
        test_loss = test_res['test_loss']
        test_acc = test_res['test_accuracy']

        plt.rcParams['figure.figsize'] = (12.8, 7.2)  # 1280 x 720 像素全局设置
        new_ticks = np.arange(1, len(train_epoch), math.ceil(len(train_epoch) / 30))

        train_acc_curve, = plt.plot(train_epoch, train_acc)
        test_acc_curve, = plt.plot(test_epoch, test_acc)
        plt.xlabel('epoch')
        plt.ylabel('accuracy')
        plt.title('accuracy-epoch')
        plt.xticks(new_ticks)
        plt.legend((train_acc_curve, test_acc_curve), ('train_accuracy', 'test_accuracy'))
        # plt.show()
        plt.savefig('../img/accuracy-epoch.png')
        plt.close()

        train_loss_curve, = plt.plot(train_epoch, train_loss)
        test_loss_curve, = plt.plot(test_epoch, test_loss)
        plt.xlabel('epoch')
        plt.ylabel('loss')
        plt.title('loss-epoch')
        plt.xticks(new_ticks)
        plt.legend((train_loss_curve, test_loss_curve), ('train_loss', 'test_loss'))
        plt.savefig('../img/loss-epoch.png')
        plt.close()

        lr_curve, = plt.plot(train_epoch, train_lr)
        plt.xlabel('epoch')
        plt.ylabel('lr')
        plt.title('lr-epoch')
        plt.xticks(new_ticks)
        plt.savefig('../img/lr-epoch.png')
        plt.close()


def start(train, test):
    train_csv_writer.writerow(['train_epoch', 'train_loss', 'train_accuracy', 'lr'])
    test_csv_writer.writerow(['test_epoch', 'test_loss', 'test_accuracy'])

    email = MyEmail('P8基于torchvision的CV迁移学习')
    step = {0: '训练环节出现问题', 1: '测试环节出现问题', 2: '分析环节出现问题'}
    step_i = -1
    try:
        print("开始训练啦")
        step_i += 1
        train()
        train_file.close()
        print('开始测试啦')
        step_i += 1
        test()
        test_file.close()
        print('开始分析啦')
        step_i += 1
        analyze()
    except:
        print('发送失败邮件啦')
        email.send_bad_email(step[step_i])
        if step_i == 0:
            train_file.close()
            test_file.close()
        if step_i == 1:
            test_file.close()
    else:
        print('发送成功邮件啦')
        email.send_good_email()

